###################################################################################

library(data.table)
library(optparse)
library(dplyr)
library(ggplot2)
library(ggsci)
library(RColorBrewer)
library(ggpubr)
library(patchwork)

###################################################################################

option_list <- list(
    make_option(c("--input_file"), type = "character") ,
    make_option(c("--ccf_file"), type = "character") ,
    make_option(c("--ccf_msi_file"), type = "character") ,
    make_option(c("--mutshare_file"), type = "character") ,
    make_option(c("--mutshare_msi_file"), type = "character") ,
    make_option(c("--images_path"), type = "character")
)

if(1!=1){
    
    input_file <- "~/20220915_gastric_multiple/dna_combinePublic/baseTable/STAD_Info.addBurden.MSI_MSS.addCNVType.tsv"
    mutshare_file <- "~/20220915_gastric_multiple/dna_combinePublic/images/mutRate/MutShare.AllPoint.tsv"
    mutshare_msi_file <- "~/20220915_gastric_multiple/dna_combinePublic/images/mutRateMSI/MutShare.AllPoint.tsv"
    ccf_file <- "~/20220915_gastric_multiple/dna_combinePublic/mutationTime/result/All_CCF_mutTime.tsv"
    ccf_msi_file <- "~/20220915_gastric_multiple/dna_combinePublic/mutationTime/result/All_CCF_mutTime.MSI.tsv"
    images_path <-"~/20220915_gastric_multiple/dna_combinePublic/images/mutBurden"

}

###########################################################################################

parseobj <- OptionParser(option_list=option_list, usage = "usage: Rscript %prog [options]")
opt <- parse_args(parseobj)
print(opt)

input_file <- opt$input_file
mutshare_msi_file <- opt$mutshare_msi_file
mutshare_file <- opt$mutshare_file
images_path <- opt$images_path
ccf_msi_file <- opt$ccf_msi_file
ccf_file <- opt$ccf_file

dir.create(images_path , recursive = T)

##############################################################################

col <- c(
  brewer.pal(9,"YlGnBu")[6],
  rgb(234,106,79,alpha=255,maxColorValue=255),
  rgb(203,24,30,alpha=255,maxColorValue=255),
  rgb(255,0,0,alpha=255,maxColorValue=255)
  )

names(col) <- c("IM" , "IGC" , "DGC" , "GC")

##############################################################################

info <- data.frame(fread(input_file))
dat_share <- data.frame(fread(mutshare_file))
dat_share_msi <- data.frame(fread(mutshare_msi_file))
dat_ccf <- data.frame(fread(ccf_file))
dat_ccf_msi <- data.frame(fread(ccf_msi_file))

dat_ccf_use <- rbind( dat_ccf , dat_ccf_msi )
dat_share_use <- rbind( dat_share , dat_share_msi )

##############################################################################
## 计算每个share的比例
dat_share_rate <- c() 
for( sample in unique(dat_share_use$Tumor) ){

    tmp <- subset( dat_share_use , Tumor == sample )
    tmp_ccf <- subset( dat_ccf_use , Sample == sample )

    tmp$vid <- paste( tmp$Chromosome , tmp$Start_Position , tmp$End_Position , sep = ":" )
    tmp_ccf$vid <- paste( tmp_ccf$Chr , tmp_ccf$Start_Position , tmp_ccf$End_Position , sep = ":" )

    tmp_use <- merge( tmp , tmp_ccf[,c("vid" , "VAF" , "CCF_adj" , "CLS")] , by = "vid" )

    share_nums <- median(subset( tmp_use , Share=="TRUE" )$CCF_adj)
    private_nums <- median(subset( tmp_use , Share=="FALSE" )$CCF_adj)
        
    dat_tmp <- data.frame( Tumor = sample , share_ccf = share_nums , private_ccf = private_nums )
    dat_share_rate <- rbind( dat_share_rate , dat_tmp )
}

##############################################################################

dat <- merge( info , dat_share_rate , by = "Tumor" )

##############################################################################

msi_sample <- unique(subset( dat , TCGA_Class %in% c("POLE" , "MSI") )$Patient)
gs_sample <- unique(subset( dat , TCGA_Class %in% c("GS") )$Patient)
cin_sample <- unique(subset( dat , TCGA_Class %in% c("CIN") )$Patient)

##############################################################################
## 分不同亚型
dat <- subset( dat , Type!="IM + IGC + DGC" )

##############################################################################

dat_use <- subset( dat , TCGA_Class == "IM" )
dat_use$Molecular <- ifelse( dat_use$Patient %in% msi_sample , "MSI" , "" )
dat_use$Molecular <- ifelse( dat_use$Patient %in% gs_sample , "GS" , dat_use$Molecular )
dat_use$Molecular <- ifelse( dat_use$Patient %in% cin_sample , "CIN" , dat_use$Molecular )

##############################################################################

trans <- function(num){
    up <- floor(log10(num))
    down <- round(num*10^(-up),2)
    text <- paste("p == ",down," %*% 10","^",up)
    return(text)
}

###########################################################################################

plotBurden <- function(dat_plot_tmp = dat_plot_tmp , type = type , share_class = share_class ){

    dat_plot_tmp$title <- type
    dat_plot_tmp$TCGA_Class <- ifelse( dat_plot_tmp$Patient %in% msi_sample , "MSI\nPOLE" , dat_plot_tmp$TCGA_Class)
    dat_plot_tmp$TCGA_Class <- ifelse( dat_plot_tmp$Patient %in% gs_sample , "GS" , dat_plot_tmp$TCGA_Class)
    dat_plot_tmp$TCGA_Class <- ifelse( dat_plot_tmp$Patient %in% cin_sample , "CIN" , dat_plot_tmp$TCGA_Class)
    dat_plot_tmp$TCGA_Class <- factor( dat_plot_tmp$TCGA_Class , levels = c("GS" , "CIN" , "MSI\nPOLE") , order = T )

    sample_num <- dat_plot_tmp %>%
    group_by( TCGA_Class ) %>%
    summarize( nums = length(unique(Patient)) )

    dat_plot_tmp <- merge( dat_plot_tmp , sample_num , by = "TCGA_Class" )
    dat_plot_tmp$TCGA_Class_num <- paste0(dat_plot_tmp$TCGA_Class , "\n" , "(" , dat_plot_tmp$nums , ")" )

    ## 一个人样本取中位数
    dat_plot_tmp_use <- dat_plot_tmp %>%
    group_by( TCGA_Class_num , Patient , title , TCGA_Class ) %>%
    summarize( share_ccf = median(share_ccf) , private_ccf = median(private_ccf) )


    if(share_class == "Share"){
        dat_plot_tmp_use$useCol <- dat_plot_tmp_use$share_ccf
        ylab <- "Median CCF of share mutations"
        ymax <- 0.11
        #label_y <- c(0.08, 0.08 , 0.09)
    }else if(share_class == "Private"){
        dat_plot_tmp_use$useCol <- dat_plot_tmp_use$private_ccf
        ylab <- "Median CCF of private mutations"
        #label_y <- c(0.45, 0.49 , 0.53)
        ymax <- 0.55
    }

    my_comparisons_1 <- list( c(1,2) , c(2,3) , c(1,3))
    col_use <- c(rgb(red=179,green=34,blue=35,alpha=255,max=255) ,
        rgb(red=247,green=184,blue=71,alpha=255,max=255) ,
        rgb(red=2,green=100,blue=190,alpha=255,max=255) 
    )

    names(col_use) <- c("GS" , "MSI\nPOLE" , "CIN")

    if(type=="Drinking"){
        ylab <- ""
        plot <- ggplot( dat_plot_tmp_use , aes( x = TCGA_Class , y = useCol , color = TCGA_Class ) ) +
            geom_line( aes( group = Patient ) , size = 0.4 , color = "gray" ) +  ## 配对样本加线
            geom_boxplot(alpha =1 , size = 0.9 , width = 0.6 , outlier.shape = NA) +
            geom_jitter(position = position_jitter(0.2) , size = 1 , alpha = 1) +
            scale_color_manual(values = col_use) +
            ylim(0,ymax) +
            facet_grid(.~title) +
            xlab(NULL) +
            ylab(ylab)+
            theme_bw() +
            #stat_compare_means(comparisons = my_comparisons_1,method = "wilcox.test" , label.y = label_y) +
            stat_compare_means(comparisons = my_comparisons_1,method = "wilcox.test") +
            theme(
                legend.position = 'none',
                legend.title = element_blank() ,
                panel.grid.major=element_blank(),
                panel.grid.minor=element_blank(),
                panel.background = element_blank(),
                panel.border = element_blank(),
                plot.title = element_text(size = 12,color="black",face='bold'),
                legend.text = element_text(size = 12,color="black",face='bold'),
                axis.text.y = element_blank(),
                axis.title.x = element_text(size = 12,color="black",face='bold'),
                axis.title.y = element_text(size = 12,color="black",face='bold'),
                axis.text.x = element_text(size = 12,color="black",face='bold') ,
                axis.ticks.length = unit(0.2, "cm") ,
                strip.text.x = element_text(size = 15, colour = "black",face='bold') ,
                axis.line = element_line(size = 0.5)
                )
    }else{
        plot <- ggplot( dat_plot_tmp_use , aes( x = TCGA_Class , y = useCol , color = TCGA_Class ) ) +
            geom_line( aes( group = Patient ) , size = 0.4 , color = "gray" ) +  ## 配对样本加线
            geom_boxplot(alpha =1 , size = 0.9 , width = 0.6 , outlier.shape = NA) +
            geom_jitter(position = position_jitter(0.2) , size = 1 , alpha = 1) +
            scale_color_manual(values = col_use) +
            ylim(0,ymax) +
            facet_grid(.~title) +
            xlab(NULL) +
            ylab(ylab)+
            theme_bw() +
            #stat_compare_means(comparisons = my_comparisons_1,method = "wilcox.test" , label.y = label_y) +
            stat_compare_means(comparisons = my_comparisons_1,method = "wilcox.test") +
            theme(
                legend.position = 'none',
                legend.title = element_blank() ,
                panel.grid.major=element_blank(),
                panel.grid.minor=element_blank(),
                panel.background = element_blank(),
                #panel.border = element_blank(),
                plot.title = element_text(size = 12,color="black",face='bold'),
                legend.text = element_text(size = 12,color="black",face='bold'),
                axis.text.y = element_text(size = 12,color="black",face='bold'),
                axis.title.x = element_text(size = 12,color="black",face='bold'),
                axis.title.y = element_text(size = 12,color="black",face='bold'),
                axis.text.x = element_text(size = 12,color="black",face='bold') ,
                axis.ticks.length = unit(0.2, "cm") ,
                strip.text.x = element_text(size = 15, colour = "black",face='bold') ,
                axis.line = element_line(size = 0.5)
                ) 
    }
    return(plot)
}

##############################################################################
## 分IM+IGC、IM+DGC分别比较
for(share_class in c("Share" , "Private")){
    type <- "All"
    dat_plot_tmp <- data.frame(dat_use)
    dat_plot_tmp <- subset( dat_plot_tmp , Class =="IM" & Type != "IM + IGC + DGC")
    p1 <- plotBurden(dat_plot_tmp = dat_plot_tmp , type = type , share_class )

    type <- "IM + IGC"
    dat_plot_tmp <- data.frame(dat_use)
    dat_plot_tmp <- subset( dat_plot_tmp , Class =="IM" & Type == "IM + IGC")
    p2 <- plotBurden(dat_plot_tmp = dat_plot_tmp , type = type , share_class )

    type <- "IM + DGC"
    dat_plot_tmp <- data.frame(dat_use)
    dat_plot_tmp <- subset( dat_plot_tmp , Class =="IM" & Type == "IM + DGC")
    p3 <- plotBurden(dat_plot_tmp = dat_plot_tmp , type = type , share_class )

    plot <- p1 + p2 + p3

    out_name <- paste0( images_path , "/mutBurden.IM.TCGA_Type.GS_CIN_MSI.MutCCFRate.",share_class,".pdf" )  
    ggsave(file=out_name,plot=plot,width=8,height=5)
}

##############################################################################
## 饮酒的影响
for(share_class in c("Share" , "Private")){
    type <- "All"
    dat_plot_tmp <- data.frame(dat_use)
    dat_plot_tmp <- subset( dat_plot_tmp , Class =="IM" & Type != "IM + IGC + DGC")
    p1 <- plotBurden(dat_plot_tmp = dat_plot_tmp , type = type , share_class )
    dat_plot_tmp$TCGA_Class <- ifelse( dat_plot_tmp$Patient %in% msi_sample , "MSI" , dat_plot_tmp$TCGA_Class)
    dat_plot_tmp$TCGA_Class <- ifelse( dat_plot_tmp$Patient %in% gs_sample , "GS" , dat_plot_tmp$TCGA_Class)
    dat_plot_tmp$TCGA_Class <- ifelse( dat_plot_tmp$Patient %in% cin_sample , "CIN" , dat_plot_tmp$TCGA_Class)
    dat_plot_tmp$TCGA_Class <- factor( dat_plot_tmp$TCGA_Class , levels = c("MSI" , "CIN" , "GS") , order = T )

    dat_plot_tmp1 <- dat_plot_tmp
    dat_plot_tmp1$type <- type

    type <- "Drinking"
    dat_plot_tmp <- data.frame(dat_use)
    dat_plot_tmp <- subset( dat_plot_tmp , Class =="IM" & Alcohol == "Drink")
    p2 <- plotBurden(dat_plot_tmp = dat_plot_tmp , type = type , share_class )
    dat_plot_tmp$TCGA_Class <- ifelse( dat_plot_tmp$Patient %in% msi_sample , "MSI" , dat_plot_tmp$TCGA_Class)
    dat_plot_tmp$TCGA_Class <- ifelse( dat_plot_tmp$Patient %in% gs_sample , "GS" , dat_plot_tmp$TCGA_Class)
    dat_plot_tmp$TCGA_Class <- ifelse( dat_plot_tmp$Patient %in% cin_sample , "CIN" , dat_plot_tmp$TCGA_Class)
    dat_plot_tmp$TCGA_Class <- factor( dat_plot_tmp$TCGA_Class , levels = c("MSI" , "CIN" , "GS") , order = T )
    dat_plot_tmp2 <- dat_plot_tmp
    dat_plot_tmp2$type <- type

    plot <- p1 + p2

    out_name <- paste0( images_path , "/mutBurden.IM.TCGA_Type.GS_CIN_MSI.MutCCFRate.",share_class,".All_Drink.pdf" )  
    ggsave(file=out_name,plot=plot,width=6/1.2,height=4.3/1.2)


    type <- "Non-Drinking"
    dat_plot_tmp <- data.frame(dat_use)
    dat_plot_tmp <- subset( dat_plot_tmp , Class =="IM" &  Alcohol == "No")
    p3 <- plotBurden(dat_plot_tmp = dat_plot_tmp , type = type , share_class )
    dat_plot_tmp$TCGA_Class <- ifelse( dat_plot_tmp$Patient %in% msi_sample , "MSI" , dat_plot_tmp$TCGA_Class)
    dat_plot_tmp$TCGA_Class <- ifelse( dat_plot_tmp$Patient %in% gs_sample , "GS" , dat_plot_tmp$TCGA_Class)
    dat_plot_tmp$TCGA_Class <- ifelse( dat_plot_tmp$Patient %in% cin_sample , "CIN" , dat_plot_tmp$TCGA_Class)
    dat_plot_tmp$TCGA_Class <- factor( dat_plot_tmp$TCGA_Class , levels = c("MSI" , "CIN" , "GS") , order = T )
    dat_plot_tmp3 <- dat_plot_tmp
    dat_plot_tmp3$type <- type

    plot <- p1 + p2 + p3

    out_name <- paste0( images_path , "/mutBurden.IM.TCGA_Type.GS_CIN_MSI.MutCCFRate.",share_class,".Non-Drink.pdf" )  
    ggsave(file=out_name,plot=plot,width=8,height=5)

    res_use <- rbind( dat_plot_tmp1 , dat_plot_tmp2 , dat_plot_tmp3 )
    res_tmp <- res_use %>%
    group_by( type , TCGA_Class ) %>%
    summarize( share_ccf = median(share_ccf) )
    out_name <- paste0( images_path , "/mutBurden.IM.TCGA_Type.GS_CIN_MSI.MutCCFRate.",share_class,".Drink.tsv" )  
    write.table( res_tmp , out_name , row.names = F , sep = "\t" , quote = F )
}
